import sys
# path to the src folder of the repository
sys.path.append('../src')

import trimesh
import numpy as np

from organ import Organ
from tissue import TissueBlock
from pipeline import Pipeline
from steps import normalize_rigid
from utils.conversions import to_array

from copy import deepcopy
from tqdm.auto import tqdm

# load source and reference organs

source_atlas = Organ(path='../data/3d-vh-m-kidney-l.glb')
hra_atlas = Organ(path='../data/3d-vh-m-kidney-r.glb')

# normalize (optional)
output = normalize_rigid(source=deepcopy(source_atlas.pointcloud), target=deepcopy(hra_atlas.pointcloud))
normalized_source_atlas = to_array(output.output['Source'])
normalized_target_atlas = to_array(output.output['Target'])

source_pc = trimesh.PointCloud(normalized_source_atlas, colors=np.tile(np.array([255, 0, 0, 1]), (len(normalized_source_atlas), 1)))
hra_pc = trimesh.PointCloud(normalized_target_atlas, colors=np.tile(np.array([0, 0, 255, 1]), (len(normalized_target_atlas), 1)))
before_scene = trimesh.Scene([source_pc, hra_pc])
before_scene.show()

trimesh.Scene([source_pc]).show()

trimesh.Scene([hra_pc]).show()

# instantiate the registration pipeline
pipeline = Pipeline(name='Base Registration', description='Base Registration', params='../configs/params.yaml')

# run registration
projections = pipeline.run(source=source_atlas, target=hra_atlas)

projected_pc = trimesh.PointCloud(projections.registration.vertices, colors=np.tile(np.array([255, 0, 0, 1]), (len(projections.registration.vertices), 1)))
hra_pc = trimesh.PointCloud(hra_atlas.vertices, colors=np.tile(np.array([0, 0, 255, 1]), (len(hra_atlas.vertices), 1)))
after_scene = trimesh.Scene([projected_pc, hra_pc])
after_scene.show()

projections.registration.show()

from sklearn.preprocessing import minmax_scale

sampled = projections.registration.simplify_quadric_decimation(0.5)

sd = trimesh.proximity.signed_distance(hra_atlas, sampled.vertices)

# rescale
sd_scaled = minmax_scale(sd, feature_range=(-1, 1), axis=0, copy=True)

# create colors
cmap = 'jet'
colors = trimesh.visual.interpolate(sd, color_map=cmap)

# rebuild the registered organ
heatmap = trimesh.Trimesh(vertices=np.array(sampled.vertices), faces=np.array(sampled.faces), vertex_colors=colors)

heatmap.show()

import seaborn as sns

sns.set_style('darkgrid')
sns.set_theme('paper')

hist = sns.histplot(data=sd_scaled).set_xlabel('Error (Signed Distance)')
fig = hist.get_figure()

!jupyter nbconvert --to html RegistrationErrorVisualization.ipynb